from Core.dataset import get_generator
from Core.Model import AlexNet
import tensorflow as tf
import os
from tensorflow.keras.optimizers import Adam
import numpy as np
from matplotlib import pyplot as plt

from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
'''
def test():
    return 1,2
a, b = test()
1 2
a, _ = test()
1  
# 说明 a, _ = test() 中 返回的是一个int类型 后续可直接运算
# a, b = test()  返回值是一个元组, 直接运算报错
'''

# 设置GPU显存自适应
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)

# 拿到处理好的数据
train_generator, validation_generator, _ = get_generator()

# 实例化模型
model = AlexNet()

# 配置训练器
model.compile(optimizer=Adam(lr=1e-4), loss='categorical_crossentropy', metrics=['acc'])

# 配置存档点和模型保存(断点续训)
checkpoint_save_path = './checkpoint/AlexNet.ckpt'
if os.path.exists(checkpoint_save_path + '.index'):
    print('--load exist model--')
    model.load_weights(checkpoint_save_path)
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                 save_weights_only=True,
                                                 save_best_only=True)
# 开始训练模型
history = model.fit_generator(train_generator,
                              steps_per_epoch=308,
                              epochs=20,
                              validation_data=validation_generator,
                              validation_steps=11,
                              callbacks=[cp_callback])
model.summary()
# 绘制结果

acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs = range(len(acc))

plt.plot(epochs, acc, 'b', label='Training accuracy')
plt.plot(epochs, val_acc, 'r', label='Validation accuracy')
plt.title('Training and validation accuracy')
plt.legend()

plt.figure()

plt.plot(epochs, loss, 'b', label='Training Loss')
plt.plot(epochs, val_loss, 'r', label='Validation Loss')
plt.title('Training and validation loss')
plt.legend()

plt.show()
